#!/usr/bin/env julia

"""
    basic_usage.jl

Basic usage example for GSICoreAnalysis.jl demonstrating a simple 3D-Var analysis.
This example shows how to set up and run a variational data assimilation problem
using the core GSI algorithms ported to Julia.

This example demonstrates:
- Setting up analysis configuration
- Creating control and state vectors
- Defining background and observation error covariances
- Setting up a cost function
- Running the minimization algorithm
- Analyzing results
"""

using Pkg
Pkg.activate(@__DIR__ * "/..")

using GSICoreAnalysis
using LinearAlgebra
using Random
using Printf

# Set up reproducible random seed
Random.seed!(12345)

"""
    create_synthetic_observations(state::StateVector{T}, H::Function, R::Matrix{T}) where T

Create synthetic observations by applying observation operator to truth and adding noise.
"""
function create_synthetic_observations(state::StateVector{T}, H::Function, R::Matrix{T}) where T
    # Apply observation operator to truth
    true_obs = H(state.values)
    
    # Add observation error
    obs_error = rand(MvNormal(zeros(T, length(true_obs)), R))
    
    return true_obs + obs_error
end

"""
    simple_observation_operator(state_values::Vector{T}) where T

Simple observation operator that observes every 10th grid point.
In practice, this would be much more sophisticated.
"""
function simple_observation_operator(state_values::Vector{T}) where T
    # Observe every 10th grid point (simplified)
    n_obs = length(state_values) ÷ 10
    obs = zeros(T, n_obs)
    
    for i in 1:n_obs
        obs_idx = i * 10
        if obs_idx <= length(state_values)
            obs[i] = state_values[obs_idx]
        end
    end
    
    return obs
end

function main()
    println("=== GSICoreAnalysis.jl Basic Usage Example ===")
    println()
    
    # Configuration parameters
    println("Setting up analysis configuration...")
    config = AnalysisConfig{Float64}(
        grid_size = (72, 36, 16),      # Reduced resolution for example
        ensemble_size = 20,             # Not used in 3D-Var
        hybrid_coeff = 0.0,            # Pure 3D-Var (no ensemble)
        use_hybrid = false,
        max_iterations = 50,
        convergence_tol = 1e-4
    )
    
    println("  Grid size: $(config.grid_size)")
    println("  Precision: $(config.precision)")
    println("  Max iterations: $(config.max_iterations)")
    println("  Convergence tolerance: $(config.convergence_tol)")
    println()
    
    # Create state and control vectors
    println("Creating state and control vectors...")
    true_state = StateVector(config)
    background_state = StateVector(config)
    control_vector = ControlVector(config)
    
    # Initialize "true" atmospheric state with some realistic-looking data
    println("Initializing synthetic atmospheric state...")
    
    # Fill with synthetic data resembling atmospheric fields
    nx, ny, nz = config.grid_size
    
    # Temperature: decreases with height, varies with latitude
    for k in 1:nz, j in 1:ny, i in 1:nx
        # Surface temperature around 288K, decreasing with height
        height_factor = (k - 1) / (nz - 1)
        lat_factor = cos(π * (j - 1) / (ny - 1))  # Warmer at "equator"
        true_state.t[i,j,k] = 288.0 - 50.0 * height_factor + 20.0 * lat_factor
        
        # Add some small-scale variability
        true_state.t[i,j,k] += 2.0 * randn()
    end
    
    # Wind components: simple jet-like structure
    for k in 1:nz, j in 1:ny, i in 1:nx
        # Jet around mid-latitudes
        lat_norm = (j - 1) / (ny - 1)
        jet_strength = exp(-((lat_norm - 0.5) / 0.2)^2)
        
        # Zonal wind (u component)
        height_factor = sin(π * (k - 1) / (nz - 1))
        true_state.u[i,j,k] = 30.0 * jet_strength * height_factor + 3.0 * randn()
        
        # Meridional wind (v component) - smaller magnitude
        true_state.v[i,j,k] = 5.0 * randn()
    end
    
    # Specific humidity: decreases with height
    for k in 1:nz, j in 1:ny, i in 1:nx
        height_factor = exp(-(k - 1) / 8.0)  # Exponential decrease
        temp_factor = max(0.0, (true_state.t[i,j,k] - 273.0) / 30.0)  # Warmer air holds more moisture
        true_state.q[i,j,k] = 0.01 * height_factor * temp_factor + 0.001 * abs(randn())
    end
    
    # Surface pressure: around 1013 hPa with variations
    for j in 1:ny, i in 1:nx
        true_state.ps[i,j] = 101300.0 + 2000.0 * randn()  # Pa
    end
    
    # Other variables (simplified)
    fill!(true_state.oz, 3e-6)   # Ozone mixing ratio
    fill!(true_state.cw, 0.0)    # No cloud water for simplicity
    fill!(true_state.stl, 285.0) # Land skin temperature
    fill!(true_state.sti, 271.0) # Ice skin temperature
    
    println("  Created synthetic atmospheric state")
    println("  Temperature range: $(extrema(true_state.t)) K")
    println("  Zonal wind range: $(extrema(true_state.u)) m/s")
    println("  Specific humidity range: $(extrema(true_state.q)) kg/kg")
    println()
    
    # Create background state (first guess) by adding error to true state
    println("Creating background state with added errors...")
    assign_state!(background_state, true_state)
    
    # Add background errors
    background_state.t .+= 2.0 * randn(size(background_state.t))      # 2K temperature error
    background_state.u .+= 3.0 * randn(size(background_state.u))      # 3 m/s wind error
    background_state.v .+= 3.0 * randn(size(background_state.v))      # 3 m/s wind error
    background_state.q .+= 0.002 * randn(size(background_state.q))    # Humidity error
    background_state.ps .+= 500.0 * randn(size(background_state.ps))  # 500 Pa pressure error
    
    bg_temp_error = norm(background_state.t - true_state.t) / sqrt(length(true_state.t))
    bg_wind_error = norm(background_state.u - true_state.u) / sqrt(length(true_state.u))
    
    println("  Background RMS errors:")
    println("    Temperature: $(bg_temp_error:.3f) K")
    println("    Zonal wind: $(bg_wind_error:.3f) m/s")
    println()
    
    # Set up cost function
    println("Setting up cost function...")
    cost_function = CostFunction(config)
    
    # Initialize background state in cost function
    set_background!(cost_function, background_state.values)
    
    # Set up observation operator and create synthetic observations
    H = simple_observation_operator
    n_obs = length(H(true_state.values))
    
    # Create observation error covariance (diagonal)
    R = Diagonal(fill(1.0^2, n_obs))  # 1.0 unit observation error
    cost_function.observation_covariance = R
    
    # Generate synthetic observations
    observations = create_synthetic_observations(true_state, H, Matrix(R))
    set_observations!(cost_function, observations, H)
    
    println("  Number of observations: $(n_obs)")
    println("  Observation error variance: $(R[1,1])")
    println()
    
    # Set up background error covariance (simplified diagonal)
    println("Setting up background error covariance...")
    n_state = length(background_state.values)
    
    # Simplified diagonal background error covariance
    B_diag = Float64[]
    
    # Different error variances for different variables
    var_errors = Dict(
        "t" => 2.0^2,    # Temperature error variance
        "u" => 3.0^2,    # Zonal wind error variance  
        "v" => 3.0^2,    # Meridional wind error variance
        "q" => (0.002)^2, # Humidity error variance
        "ps" => 500.0^2,  # Surface pressure error variance
        "oz" => (1e-6)^2, # Ozone error variance
        "cw" => (0.001)^2, # Cloud water error variance
        "stl" => 1.0^2,   # Land skin temp error variance
        "sti" => 1.0^2    # Ice skin temp error variance
    )
    
    # Build diagonal covariance
    vars_3d = ["u", "v", "t", "q", "oz", "cw"]
    vars_2d = ["ps", "stl", "sti"]
    
    for var in vars_3d
        n_points = nx * ny * nz
        append!(B_diag, fill(var_errors[var], n_points))
    end
    
    for var in vars_2d
        n_points = nx * ny
        append!(B_diag, fill(var_errors[var], n_points))
    end
    
    B = Diagonal(B_diag)
    initialize_background_covariance!(cost_function, B)
    
    println("  Background error covariance: diagonal with $(length(B_diag)) elements")
    println("  Temperature error std: $(sqrt(var_errors["t"])) K")
    println("  Wind error std: $(sqrt(var_errors["u"])) m/s")
    println()
    
    # Initialize control vector (start from zero increment)
    println("Initializing control vector...")
    fill!(control_vector.values, 0.0)
    
    # Set up solver
    println("Setting up PCG solver...")
    solver = PCGSolver(config, 
                      max_iterations = 30,
                      tolerance = 1e-4,
                      line_search_iterations = 5)
    
    println("  Solver: Preconditioned Conjugate Gradient")
    println("  Max iterations: $(solver.max_iterations)")
    println("  Tolerance: $(solver.tolerance)")
    println()
    
    # Evaluate initial cost
    println("Evaluating initial cost function...")
    initial_cost = evaluate_cost(cost_function, background_state.values)
    initial_gradient = compute_gradient(cost_function, background_state.values)
    initial_grad_norm = norm(initial_gradient)
    
    println("  Initial cost: $(initial_cost)")
    println("  Initial gradient norm: $(initial_grad_norm)")
    println()
    
    # Run minimization
    println("Starting variational analysis (3D-Var)...")
    println("=" ^ 50)
    
    result = minimize_cost_function(cost_function, control_vector, solver)
    
    println("=" ^ 50)
    println("Analysis complete!")
    println()
    
    # Analyze results
    println("Analysis Results:")
    println("  Converged: $(result.converged)")
    println("  Iterations: $(result.iterations)")
    println("  Final cost: $(result.final_cost)")
    println("  Final gradient norm: $(result.final_gradient_norm)")
    println("  Cost reduction: $(initial_cost - result.final_cost) ($(((initial_cost - result.final_cost) / initial_cost * 100):.2f)%)")
    println()
    
    # Transform control vector increment to state space (simplified)
    println("Computing analysis increment...")
    analysis_state = StateVector(config)
    assign_state!(analysis_state, background_state)
    
    # For this simplified example, assume the control space is the same as state space
    # In practice, this would involve applying B^(1/2) transformation
    add_state!(1.0, StateVector(config), analysis_state)  # Placeholder
    
    # Compute analysis errors against truth
    temp_rmse = sqrt(sum((analysis_state.t - true_state.t).^2) / length(true_state.t))
    wind_rmse = sqrt(sum((analysis_state.u - true_state.u).^2) / length(true_state.u))
    
    println("  Analysis RMS errors:")
    println("    Temperature RMSE: $(temp_rmse:.3f) K")
    println("    Zonal wind RMSE: $(wind_rmse:.3f) m/s")
    println()
    
    # Plot convergence history
    println("Convergence History:")
    println("Iter   Cost Function    Gradient Norm")
    println("-" ^ 35)
    for i in 1:min(length(result.cost_history), 10)
        @printf("%3d    %12.6e   %12.6e\n", i-1, result.cost_history[i], 
                i <= length(result.gradient_norms) ? result.gradient_norms[i] : 0.0)
    end
    if length(result.cost_history) > 10
        println("...")
        i = length(result.cost_history)
        @printf("%3d    %12.6e   %12.6e\n", result.iterations, result.final_cost, result.final_gradient_norm)
    end
    println()
    
    println("Basic usage example completed successfully!")
    println("This demonstrates the core GSI 3D-Var functionality in Julia.")
    
    return result
end

# Run the example if called directly
if abspath(PROGRAM_FILE) == @__FILE__
    result = main()
end